#include <vm_mmu.h>
#include <mm.h>
#include <mmu.h>
#include <asm/mmpage_type.h>
#include <asm/mmpage.h>
#include <asm/mmpage_pro.h>
#include <asm/base.h>
#include <sysregs.h>
#include <dump_mmpage.h>
#define NO_BLOCK_MAPPINGS BIT(0)
#define NO_CONT_MAPPINGS BIT(1)

extern char _text_boot[], _etext_boot[];
extern char _text[], _etext[];

static char gos_page_table_dir[PAGE_SIZE] __attribute__((aligned(PAGE_SIZE)));

static void alloc_init_pte(pmd_t *pmdp, unsigned long addr,
        unsigned long end, unsigned long phys,
        unsigned long prot,
        unsigned long (*alloc_pgtable)(void),
        unsigned long flags)
{
    pmd_t pmd = *pmdp;
    pte_t *ptep;

    if (pmd_none(pmd)) {
        unsigned long pte_phys;

        pte_phys = alloc_pgtable();
        set_pmd(pmdp, __pmd(pte_phys | PMD_TYPE_TABLE));
        pmd = *pmdp;
    }

    ptep = pte_offset_phys(pmdp, addr);
    do {
        set_pte(ptep, pfn_pte(phys >> PAGE_SHIFT, prot));
        phys += PAGE_SIZE;
    } while (ptep++, addr += PAGE_SIZE, addr != end);
}

static void alloc_init_pmd(pud_t *pudp, unsigned long addr,
        unsigned long end, unsigned long phys,
        unsigned long prot,
        unsigned long (*alloc_pgtable)(void),
        unsigned long flags){
    pud_t pud = *pudp;
    pmd_t *pmdp;
    unsigned long next;

    if (pud_none(pud)) {
        unsigned long pmd_phys;

        pmd_phys = alloc_pgtable();
        set_pud(pudp, __pud(pmd_phys | PUD_TYPE_TABLE));
        pud = *pudp;
    }

    pmdp = pmd_offset_phys(pudp, addr);
    do {
        next = pmd_addr_end(addr, end);

        if(((addr | next | phys) & ~SECTION_MASK) == 0 && 
                (flags & NO_BLOCK_MAPPINGS) == 0)
            pmd_set_section(pmdp, phys, prot);
        else
            alloc_init_pte(pmdp, addr, next, phys,
                    prot, alloc_pgtable, flags);
        
        phys += next -addr;
    } while (pmdp++, addr = next, addr != end);
}

static void alloc_init_pud(pgd_t *pgdp, unsigned long addr, 
        unsigned long end, unsigned long phys,
        unsigned long prot,
        unsigned long (*alloc_pgtable)(void),
        unsigned long flags)
{
    pgd_t pgd = *pgdp;
    pud_t *pudp;
    unsigned long next;

    if(pgd_none(pgd)){
        unsigned long pud_phys;

        pud_phys = alloc_pgtable();

        set_pgd(pgdp, __pgd(pud_phys | PUD_TYPE_TABLE));
        pgd = *pgdp;
    }

    pudp = pud_offset_phys(pgdp, addr);
    do {
        next = pud_addr_end(addr, end);
        alloc_init_pmd(pudp, addr, next, phys,
                prot, alloc_pgtable, flags);
        phys += next - addr;   
    } while (pudp++, addr = next, addr != end);
}

static void __create_pgd_mapping(pgd_t *pgdir, unsigned long phys,
        unsigned long virt, unsigned long size, 
        unsigned long prot,
        unsigned long (*alloc_pgtable)(void),
        unsigned long flags){
    
    printk("[MAP] VA=0x%lx->PA=0x%lx size=0x%lx prot=0x%lx\n", 
           virt, phys, size, prot);
    
    pgd_t* pgdp = pgd_offset_raw(pgdir, virt);
    unsigned long addr, end, next;

    phys &= PAGE_MASK;
    addr = virt & PAGE_MASK;
    end = PAGE_ALIGN(virt + size);

    unsigned long pgd_idx = (addr >> PGDIR_SHIFT) & (PTRS_PER_PGD - 1);
    printk("[MAP] PGD index=%lu, range=0x%lx-0x%lx\n", pgd_idx, addr, end);

    do {
        next = pgd_addr_end(addr, end);
        
        printk("[MAP] PGD[%lu] before: 0x%016lx\n", 
               pgd_idx, pgd_val(*pgdp));
        
        alloc_init_pud(pgdp, addr, next, phys, prot, alloc_pgtable, flags);
        
        printk("[MAP] PGD[%lu] after: 0x%016lx\n", 
               pgd_idx, pgd_val(*pgdp));
        
        phys += next - addr;
        pgd_idx++;

    } while (pgdp++, addr = next, addr != end);
}

/* identical_map text_boot text*/
static void create_identical_mapping(void){
    unsigned long start;
    unsigned long end;

    /*map text*/
    start = (unsigned long)_text_boot;
    end = (unsigned long)_etext;
    __create_pgd_mapping((pgd_t *)gos_page_table_dir, start, start,
            end - start, PAGE_KERNEL_ROX,
            early_pgtable_alloc,
            0);
    printk("gos_page_table_dir addr is %x \n", gos_page_table_dir);
    printk("stage 1 mapping start addr is %x \n", start);
    printk("stage 1 mapping end addr is %x \n", end);


    start = ((unsigned long)_etext);
    end = TOTAL_MEMORY;
    __create_pgd_mapping((pgd_t *)gos_page_table_dir, start, start,
            end - start, PAGE_KERNEL,
            early_pgtable_alloc,
            0);
    printk("gos_page_table_dir addr is %x \n", gos_page_table_dir);
    printk("stage 2 mapping start addr is %x \n", start);
    printk("stage 2 mapping end addr is %x \n", end);
}


static void create_mmio_mapping(){
    __create_pgd_mapping((pgd_t *)gos_page_table_dir, PBASE, PBASE,
            DEVICE_SIZE, PROT_DEVICE_nGnRnE,
            early_pgtable_alloc,
            0);
}

static void cpu_init(){
    unsigned long mair = 0;
    unsigned long tcr = 0;
    unsigned long tmp;
    unsigned long parang;

    //clear tlb
    asm("tlbi vmalle1");
    dsb(nsh);

    write_sysreg(3UL << 20, cpacr_el1);
    write_sysreg(1 << 12, mdscr_el1);

    mair = MAIR(0x00UL, MT_DEVICE_nGnRnE) |
            MAIR(0x04UL, MT_DEVICE_nGnRE) |
            MAIR(0x0cUL, MT_DEVICE_GRE) |
            MAIR(0x44UL, MT_NORMAL_NC) |
            MAIR(0xffUL, MT_NORMAL) |
            MAIR(0xbbUL, MT_NORMAL_WT);
    write_sysreg(mair, mair_el1);

    tcr = TCR_TxSZ(VA_BITS) | TCR_TG_FLAGS;

    tmp = read_sysreg(ID_AA64MMFR0_EL1);
    parang = tmp & 0xf;
    if (parang > ID_AA64MMFR0_PARANGE_48)
        parang = ID_AA64MMFR0_PARANGE_48;
    
    tcr |= parang << TCR_IPS_SHIFT;

    write_sysreg(tcr, tcr_el1);


}

static int enable_mmu(){
    unsigned long tmp;
    int tgran4;

    tmp = read_sysreg(ID_AA64MMFR0_EL1);
    tgran4 = (tmp >> ID_AA64MMFR0_TGRAN64_SHIFT) & 0xf;
    if (tgran4 != ID_AA64MMFR0_TGRAN64_SUPPORTED)
        return -1;
    
    write_sysreg(gos_page_table_dir, ttbr0_el1);

    isb();
    printk(" gos_page_table_dir is %lx in %s", gos_page_table_dir, __func__);
    write_sysreg(SCTLR_ELx_M, sctlr_el1);
    isb();

    asm("ic iallu");
    dsb(nsh);
    isb();

    return 0;
}

void gos_paging_init(){
    memset(gos_page_table_dir, 0, PAGE_SIZE);
    create_identical_mapping();
    create_mmio_mapping();
    //dump_pgtable(gos_page_table_dir, TOTAL_MEMORY);
    cpu_init();
    enable_mmu();

    printk("gos enable mmu done\n");
}

void gos_mapping(unsigned long va, unsigned long pa){
    __create_pgd_mapping((pgd_t *)gos_page_table_dir, pa, va,
            PAGE_SIZE, PAGE_KERNEL,
            early_pgtable_alloc,
            0);
}